{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# HybridGaussianConditional"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/hybrid/doc/HybridGaussianConditional.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam-develop\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `HybridGaussianConditional` represents a conditional probability distribution $P(X | M, Z)$ where the frontal variables $X$ are continuous, and the parent variables include both discrete variables $M$ and potentially other continuous variables $Z$. It inherits from `HybridGaussianFactor` and `Conditional<HybridGaussianFactor, HybridGaussianConditional>`.\n",
        "\n",
        "Essentially, it's a collection of `GaussianConditional`s, where the choice of which `GaussianConditional` applies depends on the assignment $m$ of the discrete parent variables $M$.\n",
        "$$\n",
        "P(X | M=m, Z) = P_m(X | Z)\n",
        "$$\n",
        "where $P_m(X | Z)$ is a standard `GaussianConditional`.\n",
        "\n",
        "Internally, it uses a `DecisionTree<Key, GaussianConditional::shared_ptr>` to store the different conditional components, indexed by the discrete parent `DiscreteKey`s $M$. Note that unlike `HybridGaussianFactor`, the scalar value stored alongside the factor in the base class represents the *negative log normalization constant* for that branch, ensuring each $P_m(X|Z)$ integrates to 1.\n",
        "\n",
        "These conditionals are the standard result of eliminating continuous variables when discrete variables are present in the separator during hybrid elimination."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "\n",
        "from gtsam import HybridGaussianConditional, GaussianConditional\n",
        "from gtsam.symbol_shorthand import X, D"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Can be initialized with one or more discrete parent keys and the corresponding `GaussianConditional` components."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HybridGaussianConditional P(X0 | D0):\n",
            "HybridGaussianConditional\n",
            "\n",
            " P( x0 | d0)\n",
            " Discrete Keys = (d0, 2), \n",
            " logNormalizationConstant: -0.918939\n",
            "\n",
            " Choice(d0) \n",
            " 0 Leaf p(x0)\n",
            "  R = [ 1 ]\n",
            "  d = [ 0 ]\n",
            "  mean: 1 elements\n",
            "  x0: 0\n",
            "  logNormalizationConstant: -0.918939\n",
            "  Noise model: unit (1) \n",
            "\n",
            " 1 Leaf p(x0)\n",
            "  R = [ 0.5 ]\n",
            "  d = [ 2.5 ]\n",
            "  mean: 1 elements\n",
            "  x0: 5\n",
            "  logNormalizationConstant: -2.30523\n",
            "isotropic dim=1 sigma=2\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# --- Example: Single Discrete Parent P(X0 | D0) ---\n",
        "dk0 = (D(0), 2) # Binary discrete parent D0\n",
        "\n",
        "# Gaussian conditional for mode D0=0: P(X0|D0=0) = N(0, 1)\n",
        "# R=1, d=0\n",
        "gc0 = GaussianConditional(X(0), np.zeros(1), np.eye(1), gtsam.noiseModel.Unit.Create(1))\n",
        "\n",
        "# Gaussian conditional for mode D0=1: P(X0|D0=1) = N(5, 4=2^2)\n",
        "# R=1/2=0.5, d = R*mean = 0.5*5 = 2.5\n",
        "# Need noise model with sigma=2.0\n",
        "noise1 = gtsam.noiseModel.Isotropic.Sigma(1, 2.0)\n",
        "gc1 = GaussianConditional(X(0), np.array([2.5]), np.eye(1)*0.5, noise1)\n",
        "\n",
        "hgc1 = HybridGaussianConditional(dk0, [gc0, gc1])\n",
        "print(\"HybridGaussianConditional P(X0 | D0):\")\n",
        "hgc1.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Accessing Components and Properties\n",
        "\n",
        "You can retrieve the underlying decision tree of conditionals and select a specific `GaussianConditional` based on a discrete assignment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Selected GaussianConditional for D0=0:\n",
            "GaussianConditional p(x0)\n",
            "  R = [ 1 ]\n",
            "  d = [ 0 ]\n",
            "  mean: 1 elements\n",
            "  x0: 0\n",
            "  logNormalizationConstant: -0.918939\n",
            "  Noise model: unit (1) \n",
            "\n",
            "Selected GaussianConditional for D0=1\n",
            ": (GaussianConditional p(x0)\n",
            "  R = [ 0.5 ]\n",
            "  d = [ 2.5 ]\n",
            "  mean: 1 elements\n",
            "  x0: 5\n",
            "  logNormalizationConstant: -2.30523\n",
            "isotropic dim=1 sigma=2\n",
            ", 1.3862943611198908)\n",
            "Discrete Parents (part of HybridFactor): \n",
            "d0 2\n",
            "\n",
            "Continuous Parents: []\n"
          ]
        }
      ],
      "source": [
        "# Using hgc1 from Example 1: P(X0 | D0) with modes N(0,1) and N(5,4)\n",
        "# Get the GaussianConditional for a specific assignment using choose() or ()\n",
        "assignment0 = gtsam.DiscreteValues()\n",
        "assignment0[D(0)] = 0\n",
        "selected_gc0 = hgc1.choose(assignment0) # or hgc1(assignment0)\n",
        "print(\"\\nSelected GaussianConditional for D0=0:\")\n",
        "selected_gc0.print() # Should be N(0, 1)\n",
        "\n",
        "assignment1 = gtsam.DiscreteValues()\n",
        "assignment1[D(0)] = 1\n",
        "selected_gc1 = hgc1(assignment1)\n",
        "print(f\"\\nSelected GaussianConditional for D0=1\\n: {selected_gc1}\") # Should be N(5, 4)\n",
        "\n",
        "# Access conditional properties\n",
        "print(f\"Discrete Parents (part of HybridFactor): {hgc1.discreteKeys()}\")\n",
        "print(f\"Continuous Parents: {hgc1.continuousParents()}\") # Convenience method"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluation (`logProbability`, `evaluate`)\n",
        "\n",
        "These methods evaluate the probability density $P_m(X | Z)$ corresponding to the selected discrete mode $m$. They require a `HybridValues` object containing assignments for all involved variables (frontal $X$, discrete parents $M$, continuous parents $Z$)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "P(X0=0.1 | D0=0): LogProb=-0.9239385332046728, Prob=0.39695254747701175\n",
            "P(X0=4.0 | D0=1): LogProb=-2.3364828943245635, Prob=0.0966670292007123\n",
            "Error(X0=0.1|D0=0): 0.005000000000000001 (Should be approx -log_prob0 - negLogConstant(gc0))\n",
            "  Check: 0.0050000000000000044\n"
          ]
        }
      ],
      "source": [
        "# Using hgc1: P(X0 | D0) with modes N(0,1) and N(5,4)\n",
        "\n",
        "# --- Evaluate mode 0: P(X0=0.1 | D0=0) ---\n",
        "hybrid_vals0 = gtsam.HybridValues()\n",
        "hybrid_vals0.insert(X(0), np.array([0.1])) # Frontal value\n",
        "hybrid_vals0.insert(D(0), 0)             # Discrete parent assignment\n",
        "\n",
        "log_prob0 = hgc1.logProbability(hybrid_vals0)\n",
        "prob0 = hgc1.evaluate(hybrid_vals0)\n",
        "# Expected: log density of N(0.1; 0, 1)\n",
        "print(f\"\\nP(X0=0.1 | D0=0): LogProb={log_prob0}, Prob={prob0}\")\n",
        "\n",
        "# --- Evaluate mode 1: P(X0=4.0 | D0=1) ---\n",
        "hybrid_vals1 = gtsam.HybridValues()\n",
        "hybrid_vals1.insert(X(0), np.array([4.0])) # Frontal value\n",
        "hybrid_vals1.insert(D(0), 1)             # Discrete parent assignment\n",
        "\n",
        "log_prob1 = hgc1.logProbability(hybrid_vals1)\n",
        "prob1 = hgc1.evaluate(hybrid_vals1)\n",
        "# Expected: log density of N(4.0; 5, 4)\n",
        "print(f\"P(X0=4.0 | D0=1): LogProb={log_prob1}, Prob={prob1}\")\n",
        "\n",
        "# Note: The error() method inherited from HybridGaussianFactor gives\n",
        "# the Gaussian negative log-likelihood *plus* the neg log normalization constant\n",
        "# stored in the tree.\n",
        "# error = -log P(x|m,z) - log(NormConstant)\n",
        "err0 = hgc1.error(hybrid_vals0)\n",
        "print(f\"Error(X0=0.1|D0=0): {err0} (Should be approx -log_prob0 - negLogConstant(gc0))\")\n",
        "neg_log_const_0 = gc0.negLogConstant() # Should be ~0.9189\n",
        "print(f\"  Check: {-log_prob0 - neg_log_const_0}\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
